import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE, MDS, Isomap
from umap import UMAP  # 需要单独安装：pip install umap-learn
import ast
import gc

def multiple_dim_reduction(file_path, sample_ratio=0.1):
    """使用多种降维方法可视化数据，并比较结果，分数越高颜色越深（连续渐变）"""
    # 1. 读取和预处理数据
    scores = []
    sequences = []
    
    with open(file_path, 'r') as f:
        while True:
            score_line = f.readline()
            if not score_line:
                break
            seq_line = f.readline()
            if not seq_line:
                break
                
            try:
                score = float(score_line.strip())
                sequence = ast.literal_eval(seq_line.strip())
                scores.append(score)
                sequences.append(sequence)
            except:
                continue
    
    # 转换为数组
    scores = np.array(scores)
    sequences = np.array(sequences)
    
    # 归一化分数用于着色（范围0-1）
    normalized_scores = (scores - np.min(scores)) / (np.max(scores) - np.min(scores) + 1e-10)
    for i in range(len(normalized_scores)):
        normalized_scores[i]=1-normalized_scores[i]
    # 2. 定义要比较的降维方法
    dim_reducers = {
        "PCA": PCA(n_components=2, random_state=42),
        "t-SNE": TSNE(n_components=2, random_state=42, perplexity=30),
        "UMAP": UMAP(n_components=2, random_state=42, n_neighbors=15),
        "MDS": MDS(n_components=2, random_state=42, n_init=1),
        "Isomap": Isomap(n_components=2, n_neighbors=10)
    }
    
    # 3. 对高维数据先进行PCA预处理（提高效率）
    if sequences.shape[1] > 50:
        current_dim = sequences.shape[1]
        pca_pre = PCA(n_components=64, random_state=42)
        sequences = pca_pre.fit_transform(sequences)
        print(f"PCA预处理完成，从 {current_dim} 维降至 100 维")
    
    # 4. 应用所有降维方法并可视化
    plt.figure(figsize=(20, 16))
    point_size = 15 if len(scores) > 1000 else 30
    
    # 使用红色系渐变，分数越高颜色越深
    cmap = plt.cm.Reds
    
    for i, (name, reducer) in enumerate(dim_reducers.items(), 1):
        plt.subplot(3, 2, i)
        
        # 执行降维
        print(f"正在执行 {name} 降维...")
        reduced = reducer.fit_transform(sequences)
        
        # 绘制散点图，使用连续的颜色映射
        scatter = plt.scatter(
            reduced[:, 0], reduced[:, 1],
            c=normalized_scores,  # 直接使用归一化分数作为颜色值
            cmap=cmap,
            s=point_size,
            alpha=0.7,
            edgecolors='none'
        )
        
        plt.title(f'{name} Visualization', fontsize=14)
        plt.xlabel(f'{name} Dimension 1')
        plt.ylabel(f'{name} Dimension 2')
        
        # 只在最后一个子图添加颜色条
        if i == len(dim_reducers):
            cbar = plt.colorbar(scatter, ax=plt.gca())
            cbar.set_label('Normalized Score (0.0 to 1.0)')
    
    plt.tight_layout()
    plt.savefig('dim_reduction_comparison.png', dpi=200, bbox_inches='tight')
    plt.show()
    
    return {name: reducer for name, reducer in dim_reducers.items()}

if __name__ == "__main__":
    # 根据数据规模调整采样比例
    multiple_dim_reduction('output.txt', sample_ratio=0.05)  # 对于超大规模数据，可降至1-5%
